from typing import Dict, Any, TypedDict, Annotated
from langgraph.graph import StateGraph, END
from langchain_core.messages import HumanMessage, AIMessage

class ErrorHandlingState(TypedDict):
    messages: Annotated[list, "对话消息列表"]
    result: Annotated[str, "处理结果"]
    error_info: Annotated[Dict[str, Any], "错误信息"]
    retry_count: Annotated[int, "重试次数"]
    max_retries: Annotated[int, "最大重试次数"]
    error_history: Annotated[list, "错误历史"]

class ErrorHandlingDemo:
    """错误处理演示"""
    
    def __init__(self):
        self.graph = self._build_graph()
    
    def _build_graph(self):
        """构建错误处理图"""
        
        workflow = StateGraph(ErrorHandlingState)
        
        # 添加节点
        workflow.add_node("initialize", self._initialize)
        workflow.add_node("risky_operation", self._risky_operation)
        workflow.add_node("handle_error", self._handle_error)
        workflow.add_node("retry_operation", self._retry_operation)
        workflow.add_node("fallback_operation", self._fallback_operation)
        workflow.add_node("finalize", self._finalize)
        
        # 设置入口点
        workflow.set_entry_point("initialize")
        
        # 设置条件边
        workflow.add_conditional_edges(
            "risky_operation",
            self._route_operation_result,
            {
                "success": "finalize",
                "error": "handle_error"
            }
        )
        
        workflow.add_conditional_edges(
            "handle_error",
            self._route_error_handling,
            {
                "retry": "retry_operation",
                "fallback": "fallback_operation",
                "fail": "finalize"
            }
        )
        
        # 设置其他边
        workflow.add_edge("initialize", "risky_operation")
        workflow.add_edge("retry_operation", "risky_operation")
        workflow.add_edge("fallback_operation", "finalize")
        workflow.add_edge("finalize", END)
        
        return workflow.compile()
    
    def _initialize(self, state: ErrorHandlingState) -> ErrorHandlingState:
        """初始化"""
        messages = state.get("messages", [])
        if not messages:
            messages = [HumanMessage(content="执行有风险的操作")]
        
        max_retries = state.get("max_retries", 3)
        
        return {
            "messages": messages,
            "result": "初始化完成",
            "error_info": {},
            "retry_count": 0,
            "max_retries": max_retries,
            "error_history": ["初始化开始"]
        }
    
    def _risky_operation(self, state: ErrorHandlingState) -> ErrorHandlingState:
        """有风险的操作"""
        retry_count = state["retry_count"]
        
        # 模拟有风险的操作，根据重试次数增加成功率
        import random
        success_rate = 0.3 + (retry_count * 0.2)  # 每次重试增加20%成功率
        
        if random.random() < success_rate:
            # 成功
            return {
                "messages": state["messages"],
                "result": f"操作成功 (重试次数: {retry_count})",
                "error_info": {},
                "retry_count": retry_count,
                "max_retries": state["max_retries"],
                "error_history": state["error_history"] + [f"操作成功，重试{retry_count}次"]
            }
        else:
            # 失败
            error_info = {
                "error_type": "OperationFailed",
                "error_message": f"操作失败，这是第{retry_count + 1}次尝试",
                "timestamp": "2024-01-01 10:00:00"
            }
            
            return {
                "messages": state["messages"],
                "result": "操作失败",
                "error_info": error_info,
                "retry_count": retry_count,
                "max_retries": state["max_retries"],
                "error_history": state["error_history"] + [f"操作失败: {error_info['error_message']}"]
            }
    
    def _route_operation_result(self, state: ErrorHandlingState) -> str:
        """路由操作结果"""
        return "success" if not state["error_info"] else "error"
    
    def _handle_error(self, state: ErrorHandlingState) -> ErrorHandlingState:
        """处理错误"""
        error_info = state["error_info"]
        retry_count = state["retry_count"]
        max_retries = state["max_retries"]
        
        # 更新错误历史
        error_history = state["error_history"] + [f"处理错误: {error_info['error_type']}"]
        
        return {
            "messages": state["messages"],
            "result": state["result"],
            "error_info": error_info,
            "retry_count": retry_count,
            "max_retries": max_retries,
            "error_history": error_history
        }
    
    def _route_error_handling(self, state: ErrorHandlingState) -> str:
        """路由错误处理"""
        retry_count = state["retry_count"]
        max_retries = state["max_retries"]
        
        if retry_count < max_retries:
            return "retry"
        elif retry_count == max_retries:
            return "fallback"
        else:
            return "fail"
    
    def _retry_operation(self, state: ErrorHandlingState) -> ErrorHandlingState:
        """重试操作"""
        retry_count = state["retry_count"] + 1
        
        error_history = state["error_history"] + [f"准备第{retry_count}次重试"]
        
        return {
            "messages": state["messages"],
            "result": f"准备重试 (第{retry_count}次)",
            "error_info": {},
            "retry_count": retry_count,
            "max_retries": state["max_retries"],
            "error_history": error_history
        }
    
    def _fallback_operation(self, state: ErrorHandlingState) -> ErrorHandlingState:
        """备用操作"""
        error_history = state["error_history"] + ["执行备用操作"]
        
        fallback_result = "备用操作执行成功，虽然原始操作失败，但系统仍然可用"
        
        return {
            "messages": state["messages"] + [AIMessage(content=fallback_result)],
            "result": fallback_result,
            "error_info": state["error_info"],
            "retry_count": state["retry_count"],
            "max_retries": state["max_retries"],
            "error_history": error_history
        }
    
    def _finalize(self, state: ErrorHandlingState) -> ErrorHandlingState:
        """完成处理"""
        retry_count = state["retry_count"]
        error_history = state["error_history"]
        
        if state["error_info"]:
            final_result = f"处理完成，经过{retry_count}次重试后，使用备用方案"
        else:
            final_result = f"处理完成，操作成功 (重试{retry_count}次)"
        
        error_history = error_history + ["处理完成"]
        
        return {
            "messages": state["messages"] + [AIMessage(content=final_result)],
            "result": final_result,
            "error_info": state["error_info"],
            "retry_count": retry_count,
            "max_retries": state["max_retries"],
            "error_history": error_history
        }
    
    def run(self, input_data: Dict[str, Any], config: Dict[str, Any] = None) -> Dict[str, Any]:
        """运行错误处理图"""
        max_retries = input_data.get("max_retries", 3)
        
        initial_state = {
            "messages": [HumanMessage(content=input_data.get("message", "执行有风险的操作"))],
            "result": "",
            "error_info": {},
            "retry_count": 0,
            "max_retries": max_retries,
            "error_history": []
        }
        
        result = self.graph.invoke(initial_state)
        
        return {
            "final_state": result,
            "total_retries": result["retry_count"],
            "max_retries": result["max_retries"],
            "error_info": result["error_info"],
            "error_history": result["error_history"],
            "final_result": result["result"]
        }
